from ultralytics import YOLO
import cv2
import time
import os
from PIL import Image
import Jetson.GPIO as GPIO

# 设置GPIO模式为BOARD模式
GPIO.setmode(GPIO.BOARD)

# 指定使用的GPIO引脚
button_pin = 7                                                      # 按键GPIO引脚
led_pins = [18, 23, 24, 25]                                         # 根据实际连接的引脚进行设置
for pin in led_pins:
    GPIO.setup(pin, GPIO.OUT)
GPIO.setup(button_pin, GPIO.IN, pull_up_down=GPIO.PUD_UP)


# 加载预训练的YOLOv8n模型
model = YOLO('results_runs/dataSet_good2/E10_IZ64/classify/train/weights/best.pt')

# 从摄像头捕捉图像
cap = cv2.VideoCapture(0)                                           # 0表示默认摄像头

# 设置GPIO引脚为输入模式，上拉电阻使其默认高电平
GPIO.setup(button_pin, GPIO.IN, pull_up_down=GPIO.PUD_UP)

# 去抖动处理
debounce_time = 0.3                                                 # 0.3秒的去抖动时间
last_button_state = GPIO.input(button_pin)
last_change_time = time.time()

try:
    print("垃圾识别系统开启，按下按键即可拍照识别（按下Ctrl+C结束测试）:")
    while True:
        # 检测按键是否被按下
        button_state = GPIO.input(button_pin)
        
        # 如果按键状态发生变化且经过了去抖动时间
        if button_state != last_button_state and time.time() - last_change_time > debounce_time:
            if button_state == GPIO.LOW:
                ret, frame = cap.read()                             # 读取图像帧
                if not ret:
                    break

                # 在图像上运行推理
                results = model(frame)
                
                # 处理识别结果
                top_class_id = results[0].probs.top1                # 识别序号

                # 控制LED灯
                GPIO.output(led_pins[top_class_id], GPIO.HIGH)
                time.sleep(5)                                       # 等待5秒
                GPIO.output(led_pins[top_class_id], GPIO.LOW)

            
            # 更新状态和时间
            last_button_state = button_state
            last_change_time = time.time()
            
        # 等待一段时间以避免检测到连续的按键按下
        time.sleep(0.1)
except KeyboardInterrupt:
    pass
finally:
    GPIO.cleanup()                                                  # 清理GPIO设置
    cap.release()                                                   # 释放摄像头
    cv2.destroyAllWindows()                                         # 关闭所有窗口
